from BasicLayer import *


class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.subSamples = [
            SubSample(64, 128),
            SubSample(128, 256),
            SubSample(256, 512),
            SubSample(512, 512)
        ]
        self.upSamples = [
            UpSample(1024, 256, bilinear),
            UpSample(512, 128, bilinear),
            UpSample(256, 64, bilinear),
            UpSample(128, 64, bilinear)
        ]
        self.outConv = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.subSamples[0](x1)
        x3 = self.subSamples[1](x2)
        x4 = self.subSamples[2](x3)
        x5 = self.subSamples[3](x4)
        x = self.upSamples[0](x5, x4)
        x = self.upSamples[1](x, x3)
        x = self.upSamples[2](x, x2)
        x = self.upSamples[3](x, x1)
        logistic = self.outConv(x)
        return logistic


if __name__ == '__main__':
    net = UNet(n_channels=3, n_classes=1)
    print(net)
